
from flask_socketio import SocketIO
from collections import deque
from dataclasses import dataclass
from typing import Optional
from utils.randomUtils import generate_random_by_time
from typing import Any, Optional  # 假设 SessionInfo、BaseExeHis 已定义
from .nodes.agentGraphNode import AgentGraphNode
from .nodes.baseGraphNode import  *
from .nodes.httpGraphNode import HttpGraphNode
from .nodes.sessionInfo import SessionInfo
from .nodes.toolExecuteNode import ToolExecuteNode
from .nodes.toolNode import ToolNode
from .nodes.httpInGraphNode import HttpInGraphNode
from .nodes.httpOutGraphNode import HttpOutGraphNode
from .nodes.httpAudioInGraphNode import HttpAudioInGraphNode
from .nodes.httpImageInNode import HttpImageInGraphNode
# from .nodes.sttGraphNode import STTGraphNode
from .nodes.remoteSttGraphNode import RemoteSTTGraphNode
from .nodes.MessageHistoryNode import MessageHistoryGraphNode
from .nodes.messageConcatNode import MessageConcatNode
from .nodes.manualStrNode import ManualStrGraphNode
from .nodes.agentTaskListGraphNode import AgentTaskListGraphNode
from .nodes.projectProcessGraphNode import ProjectProcessGraphNode
from .nodes.requestNode import RequestGraphNode
from .nodes.knowledgeGraphNode import KnowledgeGraphNode
from .nodes.ClearGraphNode import ClearGraphNode
from .execute.baseExeHis import BaseExeHis

@dataclass
class ExecutionTask:
    node_id: str
    session_info: SessionInfo
    exe_his: BaseExeHis
    trigger: Optional[BaseGraphNode] = None  # 若知道具体类型，可替换为实际类型（如 Trigger）
    running_id: Optional[str] = None  # 明确为字符串类型
    
class GraphNodeManager:
    def __init__(self, project):
        from .project import Project
        assert isinstance(project, Project), "project must be an instance of Project"
        self.nodesByIdDict = {}
        self.edges = None
        self.nodes = None
        self.project = project
        self.id_to_sessionInfo = {}
        self.sid_to_id = {}
        # 执行队列
        self.execution_queue = deque()
        # 记录队列中的节点ID，用于防止重复入队
        self.nodes_in_queue = set() # 判重node_id_sesseion_info_id
        # 是否正在处理队列
        self.is_processing_queue = False
        self.http_node = None

    def stop(self):
        for k,v in self.nodesByIdDict.items():
            try:
                v.stop()
            except:
                pass

        # 清空执行队列
        self.execution_queue.clear()
        self.nodes_in_queue.clear()
        self.is_processing_queue = False

    def clear_session_info(self, session_info: SessionInfo):
        session_info.clear()
        self.project.databases.get(session_info.id, {}).clear()
        print("已经清空用户历史信息")

    def get_session_info(self, id:str, socketio_instance:SocketIO, sid:str, full_url:str):
        if id not in self.id_to_sessionInfo:
            self.id_to_sessionInfo[id] = self.createSessionInfo(id, socketio_instance, sid, full_url)
        return self.id_to_sessionInfo[id]
    
    def get_session_info_by_id(self, id:str):
        return self.id_to_sessionInfo.get(id, None)

    def add_to_execution_queue(self, node_id: str, session_info: SessionInfo=None, beforerunning=False, **kwargs) -> None:
        """添加执行任务到队列"""
        # 如果是MessageHistoryGraphNode,则立即处理
        if node_id in self.nodesByIdDict and isinstance(self.nodesByIdDict[node_id], ToolExecuteNode) and session_info is None:
            self.nodesByIdDict[node_id].execute(None)
            return

        if node_id in self.nodesByIdDict and isinstance(self.nodesByIdDict[node_id], MessageHistoryGraphNode):
            self.nodesByIdDict[node_id].execute(session_info, **kwargs)# kwargs={"type":"text/image", 'role':'user/assistant'}
            return

        if node_id+"_"+session_info.id not in self.nodes_in_queue:
            task_kwargs = {}
            if kwargs.get('trigger',None):task_kwargs['trigger']=kwargs.get('trigger',None)
            if kwargs.get('running_id',None):task_kwargs['running_id']=kwargs.get('running_id',None)
            self.execution_queue.append(ExecutionTask(node_id, session_info, self.nodesByIdDict[node_id].get_new_exe_his(session_info.id ,beforerunning=beforerunning), **task_kwargs))
            self.nodes_in_queue.add(node_id+"_"+session_info.id)
            # 如果队列没有在处理，则开始处理
            if not self.is_processing_queue:
                self.process_execution_queue()

    def process_execution_queue(self) -> None:
        """处理执行队列"""
        self.is_processing_queue = True
        while self.execution_queue:
            task = self.execution_queue.popleft()
            self.nodes_in_queue.remove(task.node_id+"_"+task.session_info.id)
            try:
                # 执行节点
                self.nodesByIdDict[task.node_id].execute(task.session_info, task.exe_his, trigger=task.trigger, running_id=task.running_id)
            except Exception as e:
                print(f"Error executing node {task.node_id}: {str(e)}")
        self.is_processing_queue = False

    def deploy(self, nodes, edges, random_key:str):
        self.nodes = nodes
        self.edges = edges

        inited_http = False
        for node in nodes:
            if node['type'] == 'http': # 先创建http节点，因为http节点需要创建ws服务
                self.nodesByIdDict[node['id']] = HttpGraphNode(node['id'], node, self, random_key)
                self.http_node = self.nodesByIdDict[node['id']]
                inited_http = True
        if not inited_http:
            self.http_node = HttpGraphNode(generate_random_by_time(), {'id': generate_random_by_time(), 'type': 'http','data':{'name': 'http', 'route':generate_random_by_time(),'enWS':True,'enInterrupt':True}}, self, random_key)

        for node in nodes:
            if node['type'] == 'http':
                continue
            elif node['type'] == 'agent':
                self.nodesByIdDict[node['id']] = AgentGraphNode(node['id'], node, self, random_key)
            elif node['type']=='tool-execute':
                self.nodesByIdDict[node['id']] = ToolExecuteNode(node['id'], node, self, random_key)
            elif node['type']=='tool':
                self.nodesByIdDict[node['id']] = ToolNode(node['id'], node, self, random_key)
            elif node['type']=='http-in':
                self.nodesByIdDict[node['id']] = HttpInGraphNode(node['id'], node, self, random_key)
            elif node['type']=='http-out':
                self.nodesByIdDict[node['id']] = HttpOutGraphNode(node['id'], node, self, random_key)   
            elif node['type']=='http-audio-in':
                self.nodesByIdDict[node['id']] = HttpAudioInGraphNode(node['id'], node, self, random_key)
            elif node['type']=='http-image-in':
                self.nodesByIdDict[node['id']] = HttpImageInGraphNode(node['id'], node, self, random_key)
            elif node['type']=='message-history':
                self.nodesByIdDict[node['id']] = MessageHistoryGraphNode(node['id'], node, self, random_key)
            elif node['type']=='stt':
                # self.nodesByIdDict[node['id']] = STTGraphNode(node['id'], node, self, random_key)
                self.nodesByIdDict[node['id']] = RemoteSTTGraphNode(node['id'], node, self, random_key)
            elif node['type']=='message-concat':
                self.nodesByIdDict[node['id']] = MessageConcatNode(node['id'], node, self, random_key)
            elif node['type']=='manual-str':
                self.nodesByIdDict[node['id']] = ManualStrGraphNode(node['id'], node, self, random_key)
            elif node['type']=='agent-list-task':
                self.nodesByIdDict[node['id']] = AgentTaskListGraphNode(node['id'], node, self, random_key)
            elif node['type']=='project-process':
                self.nodesByIdDict[node['id']] = ProjectProcessGraphNode(node['id'], node, self, random_key)
            elif node['type'] == 'request':
                self.nodesByIdDict[node['id']] = RequestGraphNode(node['id'], node, self, random_key)
            elif node['type'] == 'knowledge':
                self.nodesByIdDict[node['id']] = KnowledgeGraphNode(node['id'], node, self, random_key)
            elif node['type'] == 'clear':
                self.nodesByIdDict[node['id']] = ClearGraphNode(node['id'], node, self, random_key)
            else:
                self.nodesByIdDict[node['id']] = BaseGraphNode(node['id'], node, self, random_key)

        for edge in edges:
            sourceHandle = edge['sourceHandle']
            targetHandle = edge['targetHandle']
            sourceID = edge['source']
            targetID = edge['target']
            try:
                if ('flow' in sourceHandle or 'flow' in targetHandle):
                    self.nodesByIdDict[sourceID].add_flow_next(targetID, sourceHandle=sourceHandle)
                else:
                    self.nodesByIdDict[targetID].dict['in'][targetHandle] = {'nodeId': sourceID, 'handle': sourceHandle}
            except Exception as e:
                print(f"Error adding edge {edge}: {str(e)}")

        # 基于HTTP节点属性建立WS服务处理函数
        rd = generate_random_by_time()
        for node in nodes:
            if node['type'] == 'http':
                self.add_to_execution_queue(node['id'], self.createSessionInfo(rd, None, rd, None),beforerunning=True)
            elif node['type'] == 'http-in':
                self.add_to_execution_queue(node['id'], self.createSessionInfo(rd, None, rd, None),beforerunning=True)
            elif node['type'] == 'http-audio-in':
                self.add_to_execution_queue(node['id'], self.createSessionInfo(rd, None, rd, None),beforerunning=True)
            elif node['type'] == 'http-image-in':
                self.add_to_execution_queue(node['id'], self.createSessionInfo(rd, None, rd, None),beforerunning=True)

        for node in nodes:
            if node['type'] == 'tool-execute':
                self.add_to_execution_queue(node['id'], None, beforerunning=True)
            elif node['type'] == 'project-process':
                self.add_to_execution_queue(node['id'], self.createSessionInfo(rd, None, rd, None), beforerunning=True)
    

    def createSessionInfo(self, session_id:str, ws: SocketIO, sid: str,  full_url:str):
        if(self.nodes is None) or (self.edges is None): return None
        sessionInfo = SessionInfo(session_id,ws, sid, full_url)
        for nk,nv in self.nodesByIdDict.items():
            for next_node_id in nv.get_all_flow_next_nodes():
                if next_node_id not in sessionInfo.flow_prev_nodes_executed_number:
                    sessionInfo.flow_prev_nodes_executed_number[next_node_id] = {}
                sessionInfo.flow_prev_nodes_executed_number[next_node_id][nk] = 0
        return sessionInfo
